package utils

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/des"
	"encoding/base64"
	"encoding/json"
	"fmt"
)

type AESType int

const (
	CBC AESType = iota
	ECB
	CFB
)

func (a AESType) String() string {
	switch a {
	case CBC:
		return "CBC"
	case ECB:
		return "ECB"
	case CFB:
		return "CFB"
	}
	return "UnkownType"
}

func GetRandType() AESType {
	return AESType(Rand(int(CBC), int(CFB)))
}

func AesEncode(data interface{}, key string, aesType AESType) (string, error) {
	ret := ""
	var encodeFun func(data, key []byte) ([]byte, error)

	if nil == data {
		return "", fmt.Errorf("data is nil")
	}

	if "" == key {
		return "", fmt.Errorf("key is empty")
	}

	key = MD5(key)

	if aesType.String() == "UnkownType" {
		return "", fmt.Errorf("unkown aes type")
	}

	switch aesType {
	case CBC:
		encodeFun = AesEncodeCBC
	case ECB:
		encodeFun = DesEncodeECB
	case CFB:
		encodeFun = AesEncodeCFB
	default:
		return "", fmt.Errorf("unkown aes type")
	}
	var d []byte
	var err error
	if IsString(data) {
		d = []byte(data.(string))
	} else {
		d, err = json.Marshal(data)
		if err != nil {
			return "", err
		}
	}

	encode, err := encodeFun(d, []byte(key))

	if err != nil {
		return "", err
	}

	ret = Base64Encode(encode)
	return ret, nil
}

func AesDecode(data, key string, aesType AESType) (string, error) {
	var ret string

	var fun func(data, key []byte) ([]byte, error)

	if "" == data {
		return "", fmt.Errorf("data is nil")
	}

	d, err := Base64Decode(data)
	if err != nil {
		return "", err
	}

	if "" == key {
		return "", fmt.Errorf("key is empty")
	}

	key = MD5(key)

	if aesType.String() == "UnkownType" {
		return "", fmt.Errorf("unkown aes type")
	}

	switch aesType {
	case CBC:
		fun = AesDecodeCBC
	case ECB:
		fun = DesDecodeECB
	case CFB:
		fun = AesDecodeCFB
	default:
		return "", fmt.Errorf("unkown aes type")
	}

	decode, err := fun(d, []byte(key))

	if err != nil {
		return "", err
	}

	ret = string(decode)
	return ret, nil
}

// golang 的AES加解密 （CBC/ECB/CFB 模式）
// https://blog.csdn.net/whatday/article/details/98292648
// =================== CBC ======================
func AesEncodeCBC(origData []byte, key []byte) (encrypted []byte, err error) {
	// 分组秘钥
	// NewCipher该函数限制了输入k的长度必须为16, 24或者32
	block, _ := aes.NewCipher(key)
	blockSize := block.BlockSize()                              // 获取秘钥块的长度
	origData = PKCS5Padding(origData, blockSize)                // 补全码
	blockMode := cipher.NewCBCEncrypter(block, key[:blockSize]) // 加密模式
	encrypted = make([]byte, len(origData))                     // 创建数组
	blockMode.CryptBlocks(encrypted, origData)                  // 加密
	return encrypted, nil
}
func AesDecodeCBC(encrypted []byte, key []byte) (decrypted []byte, err error) {
	block, _ := aes.NewCipher(key)                              // 分组秘钥
	blockSize := block.BlockSize()                              // 获取秘钥块的长度
	blockMode := cipher.NewCBCDecrypter(block, key[:blockSize]) // 加密模式
	decrypted = make([]byte, len(encrypted))                    // 创建数组
	blockMode.CryptBlocks(decrypted, encrypted)                 // 解密
	decrypted = PKCS5UnPadding(decrypted)                       // 去除补全码
	return decrypted, nil
}

// =================== ECB ======================

func DesEncodeECB(data, key []byte) ([]byte, error) {
	if len(key) > 8 {
		key = key[:8]
	}
	block, err := des.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("EntryptDesECB newCipher error[%v]", err)
	}
	bs := block.BlockSize()
	data = PKCS5Padding(data, bs)
	if len(data)%bs != 0 {
		return nil, fmt.Errorf("EntryptDesECB Need a multiple of the blocksize")
	}
	out := make([]byte, len(data))
	dst := out
	for len(data) > 0 {
		block.Encrypt(dst, data[:bs])
		data = data[bs:]
		dst = dst[bs:]
	}
	return out, nil
}
func DesDecodeECB(data, key []byte) ([]byte, error) {

	if len(key) > 8 {
		key = key[:8]
	}
	block, err := des.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("DecryptDES NewCipher error[%v]", err)
	}
	bs := block.BlockSize()
	if len(data)%bs != 0 {
		return nil, fmt.Errorf("DecryptDES crypto/cipher: input not full blocks")
	}
	out := make([]byte, len(data))
	dst := out
	for len(data) > 0 {
		block.Decrypt(dst, data[:bs])
		data = data[bs:]
		dst = dst[bs:]
	}
	out = PKCS5UnPadding(out)
	return out, nil
}

// =================== CFB ======================
func AesEncodeCFB(origData []byte, key []byte) (encrypted []byte, err error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	iv := key[:block.BlockSize()]

	bs := block.BlockSize()
	origData = PKCS5Padding(origData, bs)
	if len(origData)%bs != 0 {
		return nil, fmt.Errorf("EntryptDesECB Need a multiple of the blocksize")
	}

	encrypted = make([]byte, len(origData))

	stream := cipher.NewCFBEncrypter(block, iv)
	stream.XORKeyStream(encrypted, origData)
	return encrypted, nil
}
func AesDecodeCFB(encrypted []byte, key []byte) (decrypted []byte, err error) {
	block, _ := aes.NewCipher(key)
	if len(encrypted) < block.BlockSize() {
		return nil, fmt.Errorf("ciphertext too short")
	}
	iv := key[:block.BlockSize()]
	// encrypted = encrypted[aes.BlockSize:]

	stream := cipher.NewCFBDecrypter(block, iv)
	stream.XORKeyStream(encrypted, encrypted)

	encrypted = PKCS5UnPadding(encrypted)

	return encrypted, nil
}

// BASE64 ================================================

func Base64Encode(data []byte) string {
	return base64.StdEncoding.EncodeToString(data)
}

func Base64Decode(data string) ([]byte, error) {
	return base64.StdEncoding.DecodeString(data)
}

// Padding modify =======================================
func PKCS5Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padtext...)
}

func PKCS5UnPadding(origData []byte) []byte {
	length := len(origData)
	unpadding := int(origData[length-1])
	return origData[:(length - unpadding)]
}
